package Leetcode第一期;

import java.util.Random;

/**
 * @author : K k
 * @date : 22:03 2020/9/11
 * 给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。
 *
 * 注意：
 * 数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。
 *
 * 示例:
 *
 * int[] nums = new int[] {1,2,3,3,3};
 * Solution solution = new Solution(nums);
 *
 * // pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
 * solution.pick(3);
 *
 * // pick(1) 应该返回 0。因为只有nums[0]等于1。
 * solution.pick(1);
 */
public class 随机数索引_398 {

    private int[] nums;

    public 随机数索引_398(int[] nums) {
        this.nums=nums;
    }

    public int pick(int target) {
        Random r = new Random();
        int n=0;
        int index=0;
        for (int i=0;i<nums.length;i++){
            if (nums[i]==target){
                n++;
                //每一个选中概率为 k/n
                if (r.nextInt()%n==0){
                    //因为k为1，所以r.nextInt()%n==0保证是1/n的概率
                    index=i;
                }
            }
        }
        return index;
    }

    //改编 选取k个索引
    //蓄水池算法 选取第i个数的概率为 k/i  而选择并替换则是 k/i*1/k= 1/i
    public int[] pickA(int target,int k) {
        Random r = new Random();
        int n=0;
        int[] index=new int[k];
        for (int i=0;i<nums.length;i++){
            if (nums[i]==target){
                n++;
                //每一个选中概率为 k/n
                int t1=r.nextInt(n)%n;
                if (n<=k){
                    //r.nextInt()%n<=k 保证是k/n的概率 k/n是被选中的概率
                    index[n-1]=i;
                }else if(t1<k){
                    index[t1%k]=i;
                }
            }
        }
        return index;
    }

    public static void main(String[] args) {
        int[] nums = new int[] {1,2,3,3,3,3};
        随机数索引_398 solution = new 随机数索引_398(nums);
        int[] res = solution.pickA(3, 2);
        for (int i: res){
            System.out.print(i+" ");
        }
    }
}
